#ifndef PCPS_CLUSTER_HEADER
#define PCPS_CLUSTER_HEADER

#include <src/pcps.h>
#include <src/operator.h>

namespace pcps {

  //-------------------------------------------------------------------------------
  // pcps::Cluster -- A class for storing and incrementing the occupations of
  //                  a cluster of sites.
  //
  //                  Note that the pcps::Cluster class uses the default copy
  //                  constructor and assignment operator, so all copies are deep.
  //-------------------------------------------------------------------------------

  class Cluster {

    // private data
    private:
      int _size; // the number of sites in the cluster
      int _configs_per_site; // the maximum number of occupations a site can have
      bool _finished; // whether or not all configurations of the sites have been looped over
      std::vector<int> _site_lookup; // a table that converts the site's number to its position in the arrays below
      std::vector<int> _sites; // a list of the sites whose occupations are being managed
      std::vector<int> _offsets; // required offsets between the sites' bra and ket occupations (bra occ = ket occ + offset)
      std::vector<int> _n_allowed_occs; // for each site, the number of allowed ket occupations
      std::vector<int> _occ_nums; // for each site, the number of the site's occupation (this is not the site's occupation)
      std::vector< std::vector<int> > _allowed_occs; // for each site, the allowed occupations of the site in the wavefunction ket
      std::vector<int> _ket_occs; // the current occupations of the cluster's sites in the ket (-1 means the site is not in the cluster)

    // constructors
    public:
      Cluster() : _configs_per_site(0), _size(0), _finished(true) {}
      Cluster(const int n_lattice_sites,
              const int configs_per_site);
      Cluster(const int n_lattice_sites,
              const int configs_per_site,
              const std::set<int> & sites);
      Cluster(const int n_lattice_sites,
              const int configs_per_site,
              const std::vector<int> & sites);
      template <class SCALAR>
      Cluster(const pcps::Input & userinp,
              const pcps::Wavefunction<SCALAR> & wfn,
              const typename pcps::OpBase<SCALAR>::const_ptr & op,
              const pcps::Cluster & fixed_bra_occs = pcps::Cluster(),
              const pcps::Cluster & fixed_ket_occs = pcps::Cluster(),
              const std::vector<int> & pert_corr_pos = std::vector<int>());

    // public member functions
    public:

      // functions to return data
      int size() const { return _size; }
      const std::vector<int> & sites() const { return _sites; }
      const std::vector<int> & offsets() const { return _offsets; }
      const std::vector<int> & occ_nums() const { return _occ_nums; }
      const std::vector<int> & n_allowed_occs() const { return _n_allowed_occs; }
      const std::vector< std::vector<int> > & allowed_occs() const { return _allowed_occs; }
      const std::vector<int> & ket_occs() const { return _ket_occs; }
      std::vector<int> & ket_occs() { return _ket_occs; }

      // function to check whether a site is present in the cluster
      bool has_site(const int site) const {
        assert( site >= 0 && site < _site_lookup.size() );
        return (_site_lookup[site] >= 0);
      }

      // function to return the total occupation of the cluster's sites
      int total_occ() const {
        int retval = 0;
        for (int i = 0; i < _size; i++)
          retval += _ket_occs[_sites[i]];
        return retval;
      }

      // function to set a site's occupation
      void set_occ(const int site, const int occ) {
        assert( site >= 0 && site < _site_lookup.size() );
        const int site_pos = _site_lookup[site];
        assert( site_pos >= 0 && site_pos < _occ_nums.size() );
        int i;
        for (i = 0; i < _n_allowed_occs[site_pos]; i++)
          if (_allowed_occs[site_pos][i] == occ) {
            _occ_nums[site_pos] = i;
            _ket_occs[site] = occ;
            break;
          }
        if (i == _n_allowed_occs[site_pos])
          throw pcps::Exception("impossible occupation in pcps::Cluster::set_occ");
      }

      // function to fix a site's occupation
      void fix_occ(const int site, const int occ) {
        assert( site >= 0 && site < _site_lookup.size() );
        const int site_pos = _site_lookup[site];
        assert( site_pos >= 0 && site_pos < _occ_nums.size() );
        int i;
        for (i = 0; i < _n_allowed_occs[site_pos]; i++)
          if (_allowed_occs[site_pos][i] == occ)
            break;
        if (i == _n_allowed_occs[site_pos])
          throw pcps::Exception("impossible occupation in pcps::Cluster::fix_occ");
        _allowed_occs[site_pos][0] = occ;
        _n_allowed_occs[site_pos] = 1;
        _occ_nums[site_pos] = 0;
        _ket_occs[site] = occ;
      }

      // function to increment to the next occupation using the prefix ++ operator
      void operator++() {
        int i;
        bool to_break = false;
        for (i = _size - 1; i >= 0; i--) {
          int & occ = _occ_nums[i];
          if ( ++occ == _n_allowed_occs[i])
            occ = 0;
          else
            to_break = true;
          _ket_occs[_sites[i]] = _allowed_occs[i][occ];
          if (to_break) break;
        }
        if (i < 0)
          _finished = true;
      }

      // function to increment using the postfix ++ operator
      void operator++(int) { ++(*this); }

      // function to check if we have finished looping over all occupations
      bool finished() const { return _finished; }

      // function to get the ket occupation of a site
      int ket_occ(const int site) const {
        //assert( site >= 0 && site < _site_lookup.size() );
        //const int site_pos = _site_lookup[site];
        //assert( site_pos >= 0 && site_pos < _occ_nums.size() );
        //assert( _occ_nums[site_pos] < _n_allowed_occs[site_pos] );
        //return _allowed_occs[site_pos][_occ_nums[site_pos]];
        assert( site >= 0 && site < _ket_occs.size() );
        assert( _site_lookup[site] >= 0 && _site_lookup[site] < _occ_nums.size() );
        return _ket_occs[site];
      }

      // function to get the bra occupation of a site
      int bra_occ(const int site) const {
        //assert( site >= 0 && site < _site_lookup.size() );
        //const int site_pos = _site_lookup[site];
        //assert( site_pos >= 0 && site_pos < _occ_nums.size() );
        //assert( _occ_nums[site_pos] < _n_allowed_occs[site_pos] );
        //return _allowed_occs[site_pos][_occ_nums[site_pos]] + _offsets[site_pos];
        assert( site >= 0 && site < _site_lookup.size() && site < _ket_occs.size() );
        assert( _site_lookup[site] >= 0 && _site_lookup[site] < _offsets.size() );
        return _ket_occs[site] + _offsets[_site_lookup[site]];
      }

      // function to compute a compound index from the sites' occupations
      int cmpd() const {
        int retval = 0;
        for (int i = 0; i < _size; i++)
          retval = retval * _configs_per_site + _ket_occs[_sites[i]];
          //retval = retval * _configs_per_site + _allowed_occs[i][_occ_nums[i]];
        return retval;
      }

      //// function to compare two occupations
      //bool operator==(const Cluster &other) const {
      //  return (   _size == other._size
      //          && _configs_per_site == other._configs_per_site
      //          && _finished == other._finished
      //          && _site_lookup == other._site_lookup
      //          && _sites == other._sites
      //          && _offsets == other._offsets
      //          && _n_allowed_occs == other._n_allowed_occs
      //          && _occ_nums == other._occ_nums
      //          && _allowed_occs == other._allowed_occs );
      //}

      // function to reset the occupations to their initial values
      void reset_occs() {
        _finished = ( this->nconfigs() == 0 );
        for (int i = 0; i < _size; i++) {
          _occ_nums[i] = 0;
          if (_n_allowed_occs[i] > 0)
            _ket_occs[_sites[i]] = _allowed_occs[i][0];
        }
      }

      // function to find the total number of configurations of the sites
      int nconfigs() const {
        assert(_size >= 0);
        int retval = std::min(1, _size);
        for (int i = 0; i < _size; i++)
          retval *= _n_allowed_occs[i];
        return retval;
      }

      // function to split a cluster into smaller clusters
      void split(const int max_configs, std::list<pcps::Cluster> & subclusters) const;

      // functions to send and receive a cluster object for MPI
      void send(const MPI::Comm & comm, const int dest, const int tag) const;
      void recv(const MPI::Comm & comm, const int source, const int tag);

    // private member functions
    private:
      
      // function to add a site
      void add_site(const int site, const int offset);

  };

}

#endif
